import cv2
import utils

TSRW, TSRH, TSRX, TSRY = 1200, 600, 360, 220


def draw_something(frame):
    # 画框， 注意：若画的线太细经过缩放线可能消失
    frame = cv2.rectangle(frame, (TSRX, TSRY), (TSRX + TSRW, TSRY + TSRH), (0, 0, 255), 2)
    cv2.line(frame, (0, 900), (1920, 900), (0, 0, 255), 2)
    cv2.line(frame, (960, 700), (600, 900), (0, 0, 255), 2)
    cv2.line(frame, (960, 700), (1320, 900), (0, 0, 255), 2)
    # 图片缩放
    # frame = cv2.resize(frame, (960, 720))
    return frame


path_video = r"/media/fang/Samsung USB/G410D-0623/G410D-0623/CARDV/Movie/20210101032451_000016A.MP4"  # 要播放的视频路径
OutPutImgPath = r'/media/fang/Samsung USB/temp'  # 要保存图片的目录
Start = 170  # 从第Start秒播放视频

video = cv2.VideoCapture(path_video)
fps = video.get(cv2.CAP_PROP_FPS)
StarFrame = fps * Start
video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
NameNum = 300
BinNum = 0
auto_model = False
success, frame = video.read()
cnt_s = 0
while success:
    # 图像处理部分，可在函数内画框、画线、缩放
    frame = draw_something(frame)
    cv2.imshow('frame', frame)
    if auto_model:
        key = cv2.waitKey(int(1000 / int(fps * auto_model))) & 0xFF
    else:
        key = cv2.waitKey(0) & 0xFF
    # if key != -1 and key != 255:
    #     print('get key :', chr(key), 'value :', key)
    if ord('2') == key:  # 保存处理后的图片
        NameNum += 1
        if cv2.imwrite(OutPutImgPath + '/V%05d.jpg' % NameNum, frame):
            print('Write %s successful' % OutPutImgPath + '/V%05d.jpg' % NameNum)  # 图片保存成功
        else:
            print('Write %s failure!' % OutPutImgPath + '/V%05d.jpg' % NameNum)  # 图片保存失败，或因为是中文路径
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
    elif ord('w') == key:  # 保存原图
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
        success, frame = video.read()
        BinNum += 1
        # iframemg = frame[TSRY: TSRY + TSRH, TSRX: TSRX + TSRW]
        # frame = cv2.resize(frame, (600, 300))
        # utils.bgr2yuv420bin(OutPutImgPath + '\\%05d.bin' % BinNum, frame)
        NameNum += 1
        name_jpg = OutPutImgPath + '/%d.jpg' % NameNum
        if cv2.imwrite(name_jpg, frame):
            print('Write %s successful' % name_jpg)  # 图片保存成功
        else:
            print('Write %s failure!' % name_jpg)  # 图片保存失败，或因为是中文路径
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
    elif ord('a') == key:  # 按下a返回上一帧
        StarFrame -= 1
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
    elif ord('d') == key or ord('D') == key:  # 按d切换到下一帧
        StarFrame += 1
    elif ord('s') == key:  # 按s开启/关闭自动播放模式
        cnt_s += 2
        if cnt_s == 2:
            cnt_s = 1
        cnt_s = cnt_s % 5
        auto_model = cnt_s
        StarFrame += 1
    elif ord('e') == key:  # 按e将视频快进2秒
        StarFrame += 2 * fps
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
    elif ord('q') == key:  # 按q将视频后退两秒
        StarFrame -= 2 * fps
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)
    elif auto_model:
        StarFrame += auto_model
    elif -1 != key:
        video.set(cv2.CAP_PROP_POS_FRAMES, StarFrame)

    success, frame = video.read()  # 获取下一帧
video.release()
